'''
Created on May 31, 2011

@author: Kathy Benitez

Contains the main functionality for setting up population statistic databases, using the files 
provided by the U.S. Census (2000).  See accompanying README for further instructions
'''
import os, csv, re, argparse
import popStats, settings

FILE_RE = re.compile('^dc_dec_2000_sf1_u_data\d.txt')
ZIP_RE = re.compile('(\w{5}) 5-Digit ZCTA')
COUNTY_RE = re.compile('([-\w\s\']+),\s*([-\w\s\'s]+)')
DESC_RE = re.compile('People who are ([\w\s]+):\s*(\w+)\s*(;\s*([\w\s]+))?')

def processFieldnames(popStats, pathToCensusFiles, addToDB=True):
    """
    Given a directory location, get the fieldname information into the database
    """
    summaryFields = []
    for root, dirs, files in os.walk(pathToCensusFiles):
        for file in files:
            if(FILE_RE.match(file)):
                f = open(os.path.join(root,dirs,file), 'rb')
                reader = csv.DictReader(f, delimiter="|")
                desc = reader.next()
                fn = getFieldnames(desc)
                summaryFields.append([f['fieldname'] for f in fn if f['minage'] == 0 and f['maxage'] == 120])
                if addToDB:
                    popStats.description.insert().execute(*fn)
    popStats.summaryFields = summaryFields
    
def processZipPops(popStats, pathToCensusFiles):
    """
    Given a directory location for the census files, insert the data into the
    database
    """
    for root, dirs, files in os.walk(pathToCensusFiles):
        for file in files:
            if(FILE_RE.match(file)):
                f = open(os.path.join(root,dirs, file), 'rb')
                reader = csv.DictReader(f, delimiter="|")
                reader.next()    
                for geography in reader:
                    pops = getZipPopulations(geography)
                    pops2 = [p for p in pops if worthKeeping(popStats.summaryFields, p)]
                    popStats.records.insert().execute(*pops2)

def processCountyPops(popStats, pathToCensusFiles):
    """
    Given a directory location for the census files, insert the data into the
    database
    """
    for root, dirs, files in os.walk(pathToCensusFiles):
        for file in files:
            if(FILE_RE.match(file)):
                f = open(os.path.join(root,dirs, file), 'rb')
                reader = csv.DictReader(f, delimiter="|")
                reader.next()    
                for geography in reader:
                    pops = getCountyPopulations(geography)
                    pops2 = [p for p in pops if worthKeeping(popStats.summaryFields, p)]
                    popStats.records.insert().execute(*pops2)
                    
def getFieldnames(dict):
    """
    We match each of the field descriptions in the second
    line of the census files with a regular expression,
    enabling us to parse the race, sex, and age (or ages)
    encoded in that field.
    
    Example:  "People who are White alone, Female, 35 years"
    translates to a dictionary:  {"race":"White alone", "sex":"Female",
    "minage":35, "maxage":35}
    """
    fieldnames = []
    for field in dict:
        if field[:3]=="PCT":
            try:
                d = DESC_RE.search(dict[field])
                ages = ageRange(d.group(4))
                fieldnames.append({"fieldname":field,
                               "race":d.group(1),
                               "sex":d.group(2),
                               "minage":ages[0],
                               "maxage":ages[1]})
            except AttributeError as e:
                print e, "Unable to parse ", dict[field]
                continue
            except IndexError as e:
                print "No such group", e
                continue
    return fieldnames

def getZipPopulations(dict):
    """
    After the first two rows in the census file,
    each line contains a locality (in this case, a
    zip code), and populations for each of many fields.
    We parse out the zip code, then create an entry
    for each field to be added to the database.
    """
    zip_entries = []
    try:
        zip5 = ZIP_RE.search(dict["GEO_NAME"]).group(1)
    except AttributeError as e:
        print "Unable to parse", dict["GEO_NAME"], e
        return {}
    for field in dict:
        if field[:3]=="PCT":
            if dict[field]:
                pop = int(dict[field])
            else:
                pop = 0
            zip_entries.append({"zip5":zip5,
                                "fieldname_id":field,
                                "population":pop})  
    return zip_entries

def getCountyPopulations(dict):
    """
    After the first two rows in the census file,
    each line contains a locality (in this case, a
    zip code), and populations for each of many fields.
    We parse out the county and state, then create an entry
    for each field to be added to the database.
    """
    county_entries = []
    try:
        match = COUNTY_RE.search(dict["GEO_NAME"])
        county = match.group(1)
        state = match.group(2)
    except AttributeError as e:
        print "Unable to parse", dict["GEO_NAME"], e
        return {}
    for field in dict:
        if field[:3]=="PCT":
            if dict[field]:
                pop = int(dict[field])
            else:
                pop = 0
            county_entries.append({"county":county,
                                "state":state,
                                "fieldname_id":field,
                                "population":pop})  
    return county_entries

def worthKeeping(summaryFields, pop):
    """ We only want fields with some population, which are not summary fields (which
    duplicate data from the detail fields """
    return pop['fieldname_id'] not in summaryFields and pop['population'] > 0

def ageRange(inStr):
    """Return a tuple (min, max) for a string describing an age range
    Possible values: Under 1 year, 1 year, [2-99] years, 
                     [100 to 104, 105 to 109] years, 110 years and over, empty string
    >>> ageRange("")
    (0, 120)
    >>> ageRange("Under 1 year")
    (0, 0)
    >>> ageRange("1 year")
    (1, 1)
    >>> ageRange("2 years")
    (2, 2)
    >>> ageRange("2 year")
    (2, 2)
    >>> ageRange("58 years")
    (58, 58)
    >>> ageRange(" 58 years    ")
    (58, 58)
    >>> ageRange("58    years")
    (58, 58)
    >>> ageRange("100 to 104 years")
    (100, 104)
    >>> ageRange("110 years and over")
    (110, 120)
    >>> ageRange("Lorem ipsum 24")
    Traceback (most recent call last):
    ...
    TypeError: Lorem ipsum 24
    """
    if not inStr:
        return (0, 120)
    str = inStr.strip()
    if re.match('^Under\s1\syear$', str):
        l = re.match('^Under\s1\syear$', str)
        min = 0
        max = 0
    elif re.match('^(1)\s+year$', str):
        l = re.match('^(1)\s+year$', str)
        min = l.group(1)
        max = l.group(1)
    elif re.match('^(\d+)\s+year(s)?$', str):
        l = re.match('^(\d+)\s+year(s)?$', str)
        min = l.group(1)
        max = l.group(1)
    elif re.match('^(\d+)\s+to\s+(\d+)\s+years$', str):
        l = re.match('^(\d+)\s+to\s+(\d+)\s+years$', str)
        min = l.group(1)
        max = l.group(2)
    elif re.match('^(\d+)\syears\s+and\s+over$', str):
        l = re.match('^(\d+)\syears\s+and\s+over$', str)
        min = l.group(1)
        max = 120
    else:
        raise TypeError(inStr)
    return (int(min), int(max))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Set up database files in the form expected by the vdart software.')
    parser.add_argument('--e', dest='engine', default = settings.engine)
    parser.add_argument('--u', dest='uname', default = settings.uname)
    parser.add_argument('--p', dest='passwd', default = settings.passwd)
    parser.add_argument('--h', dest='hostAndPort', default = settings.hostAndPort)
    parser.add_argument('--d', dest='database', default = settings.database)
    parser.add_argument('--z', dest='pathToCensusZipFiles', default = settings.pathToCensusZipFiles)
    parser.add_argument('--c', dest='pathToCensusCountyFiles', default = settings.pathToCensusCountyFiles)
    args = parser.parse_args()
    
    #open database. Check for tables.  If they exist, check for data integrity
    if args.pathToCensusZipFiles:
        ps = popStats.PopStatsZip(args.engine, args.uname, args.passwd, 
                                  args.hostAndPort, args.database, echo=True)
        if not ps.description.exists():
            ps.description.create()
            processFieldnames(ps,args.pathToCensusZipFiles)
        if not ps.records.exists():
            if not ps.summaryFields:
                processFieldnames(ps, args.pathtoCensusZipFiles, addToDB=False)
            ps.records.create()
            processZipPops(ps, args.pathToCensusZipFiles)
        okay = ps.testIntegrity()
        if okay:
            print "Your Zip tables are correctly set up.  You should be able to perform analysis using ZIP codes."
        else:
            print "There is a problem with your Zip table setup."
    if args.pathToCensusCountyFiles:
        ps = popStats.PopStatsCounty(args.engine, args.uname, args.passwd, 
                                  args.hostAndPort, args.database, echo=True)
        if not ps.description.exists():
            ps.description.create()
            processFieldnames(ps, args.pathToCensusCountyFiles)
        if not ps.records.exists():
            if not ps.summaryFields:
                processFieldnames(ps, args.pathToCensusCountyFiles, addToDB=False)
            ps.records.create()
            processCountyPops(ps, args.pathToCensusZipFiles)
        okay = ps.testIntegrity()
        if okay:
            print "Your County tables are correctly set up.  You should be able to perform analysis using Counties."
        else:
            print "There is a problem with your County table setup."